#include "../DWMRICore/DTI.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/VectorField.h"

#include <teem/ten.h>

#include <omp.h>

#define _USE_MATH_DEFINES
#include <math.h>
#include <time.h>
#include <stdio.h>

char *pWMMaskPathname = NULL;
char *pDTIPathname = NULL;
char *pVECPathname = NULL;

/* MLS fiber tracking parameters */
float MLSSize = 4.0f;
CPoint3F *pSamplePoints = NULL;
float *pSamplePointsWeight = NULL;

CDTI *pDTI = NULL;					/* DTI */
CScalar *pWMMASK = NULL;			/* white matter */

int scale = 2;

void PrintUsage()
{
	printf("Usage: ");
	printf("DTI2VEC -dti <DTI> -wm_mask <WMMask> -vec <VectorField>\n");
	printf("\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-dti") == 0) {
			pDTIPathname = argv[++i];
		} else if (strcmp(argv[i], "-wm_mask") == 0) {
			pWMMaskPathname = argv[++i];
		} else if (strcmp(argv[i], "-vec") == 0) {
			pVECPathname = argv[++i];
		} else {
            printf("unknown %s\n", argv[i]);
			return 1;		// error
		}
	}

	if (pDTIPathname == NULL || pWMMaskPathname == NULL || pVECPathname == NULL)
		return -2;

	return 0;
}


void CreateMLSSamplePoints()
{
	SafeDeleteArray(pSamplePoints);
	SafeDeleteArray(pSamplePointsWeight);
	pSamplePoints = new CPoint3F[SAMPLE_POINT_NUM];
	pSamplePointsWeight = new float[SAMPLE_POINT_NUM];

	pSamplePoints[0].m_x = 0.0f;
	pSamplePoints[0].m_y = 0.0f;
	pSamplePoints[0].m_z = 0.0f;
	pSamplePointsWeight[0] = 1.0f;

	clock_t seed = clock();
	srand(seed);

	for (int i = 1; i < SAMPLE_POINT_NUM; ++i) {
		bool flag = false;
		CPoint3F pt;
		while(!flag) {
			float a = (float)rand() / (float)RAND_MAX * 2.0f * (float)M_PI;
			float b = (float)rand() / (float)RAND_MAX * (float)M_PI;
			float r = (float)rand() / (float)RAND_MAX;

			pt.m_x = r * sinf(a) * sinf(b);
			pt.m_y = r * sinf(a) * cosf(b);
			pt.m_z = r * cosf(a);

			flag = true;
			for (int j = 0; j < i; ++j) {
				float dis = Distance(pSamplePoints[j], pt);
				if (dis < 0.25f) {
					flag = false;
					break;
				}
			}
		}

		pSamplePoints[i] = pt;
		pSamplePointsWeight[i] = InnerProduct(pSamplePoints[i], pSamplePoints[i]) * 4.0f;
		pSamplePointsWeight[i] = expf(pSamplePointsWeight[i] / -2.5f);
	}
}


void DTIMLSFilter(const CPoint3F pos, float *lastDTI, float *filteredDTI)
{
	float m[9];
	TEN_T2M(m, lastDTI);

	CPoint3F samples[SAMPLE_POINT_NUM];
	float length[SAMPLE_POINT_NUM];
	float max_length = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = pSamplePoints[i].m_x * m[0] + pSamplePoints[i].m_y * m[1] + pSamplePoints[i].m_z * m[2];
		samples[i].m_y = pSamplePoints[i].m_x * m[3] + pSamplePoints[i].m_y * m[4] + pSamplePoints[i].m_z * m[5];
		samples[i].m_z = pSamplePoints[i].m_x * m[6] + pSamplePoints[i].m_y * m[7] + pSamplePoints[i].m_z * m[8];
		length[i] = sqrtf(InnerProduct(samples[i], samples[i]));
		if (length[i] > max_length)
			max_length = length[i];
	}

	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		samples[i].m_x = samples[i].m_x / max_length * MLSSize + pos.m_x;
		samples[i].m_y = samples[i].m_y / max_length * MLSSize + pos.m_y;
		samples[i].m_z = samples[i].m_z / max_length * MLSSize + pos.m_z;
	}

	memset(filteredDTI, 0, sizeof(float)*7);
	float temp_DTI[7];
	float weight = 0.0f;
	for (uint i = 0; i < SAMPLE_POINT_NUM; ++i) {
		if (pWMMASK->GetScalar(samples[i]) < 0.5f) {
			continue;
		}
		pDTI->GetDTI(samples[i], temp_DTI);

		for (uint j = 0; j < 7; ++j) {
			filteredDTI[j] += pSamplePointsWeight[i] * temp_DTI[j];
		}
		weight += pSamplePointsWeight[i];
	}

	for (uint j = 0; j < 7; ++j) {
		filteredDTI[j] = filteredDTI[j] / weight;
	}
}

CVector3F DTI2VEC(const CPoint3F pos)
{
	if (pWMMASK->GetScalar(pos) < 0.5f) {
		return CVector3F(0.0f, 0.0f, 0.0f);
	}

	float initDTI[7] = {1, 1, 0, 0, 1, 0, 1};
	float filteredDTI[7];
	DTIMLSFilter(pos, initDTI, filteredDTI);
	CVector3F dir;
	pDTI->DTI2Direction(filteredDTI, &dir);
	return dir;
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}

	pDTI = new CDTI();
	pDTI->OpenDTIFile(pDTIPathname);
	pWMMASK = new CScalar();
	pWMMASK->OpenScalarFile(pWMMaskPathname);

	CreateMLSSamplePoints();

	uint w = pDTI->GetWidth();
	uint h = pDTI->GetHeight();
	uint d = pDTI->GetDepth();

	int sw = (int)w * scale;
	int sh = (int)h * scale;
	int sd = (int)d * scale;

	CVector3F *vector = new CVector3F[sw*sh*sd];
	for (int z = 0; z < sd; ++z) {
		omp_set_num_threads(16);
#pragma omp parallel for
		for (int y = 0; y < sd; ++y) {
			for (int x = 0; x < sw; ++x) {
				CPoint3F pos;
				pos.m_x = (float)x/(float)scale;
				pos.m_y = (float)y/(float)scale;
				pos.m_z = (float)z/(float)scale;
				int index = (z * sh + y) * sw + x;
				vector[index] = DTI2VEC(pos);
			}
		}
		float ff = float(z+1) / (float)sd * 100.0f;
		printf("\r%.2f%% done.", ff);
	}

	CVectorField *pVectorField = new CVectorField();
	pVectorField->CreateVectorField(sw, sh, sd);
	for (int i = 0; i < sw*sh*sd; ++i) {
		pVectorField->SetVector(i, vector[i]);
	}
	pVectorField->SaveVectorFieldFile(pVECPathname);

	SafeDelete(pDTI);
	SafeDelete(pWMMASK);
	SafeDelete(pVectorField);

	SafeDeleteArray(vector);

	return 0;
}